from langchain.document_loaders.unstructured import UnstructuredFileLoader
from typing import List, Dict, Optional
import tqdm
import json

from config.crypto_config import *


class RapidOCRDocLoader(UnstructuredFileLoader):
    def __init__(self, file_path: str, custom_doc_key_extraction: Optional[Dict[str, str]] = None, **kwargs):
        super().__init__(file_path, **kwargs)
        self.custom_doc_key_extraction = custom_doc_key_extraction
        
        # 实例变量，确保每次调用都是独立的数据
        self.CRYPTO_EVAL_INFORMATION = {
            "KEY": "",
            "AUTH": "",
            "ACCESS_CONTROL": "",
            "KEY_DATA_TRANSPORT_CONFIDENTIALITY_AND_INTEGRITY": "",
            "KEY_DATA_STORE_CONFIDENTIALITY": "",
            "KEY_DATA_STORE_INTEGRITY": "",
            "KEY_DATA_SAFE_MARK_INTEGRITY": "",
            "NON-REPUDIATION": "",
        }
        
        self.DOC_MD_CONTENT = []
        self.DOC_MD_TITLE = []

    def getDockKeyParagraph(self):
        return self.CRYPTO_EVAL_INFORMATION

    def getDocMDContent(self):
        return self.DOC_MD_CONTENT

    def getDocMDTitle(self):
        return self.DOC_MD_TITLE

    def convert_table_to_md(self, table):
        if not table:
            return None
        md_lines = []
        header = "| " + " | ".join(table[0]) + " |"
        md_lines.append(header)
        separator = "| " + " | ".join(["---" for _ in table[0]]) + " |"
        md_lines.append(separator)

        # 添加表格内容行，确保每行都是新的一行
        for row in table[1:]:
            row_str = "| " + " | ".join(row) + " |"
            md_lines.append(row_str)

        # 表格后添加空行
        md_lines.append("")

        # 清理多余的空行（连续空行只保留一个）
        cleaned_lines = []
        prev_empty = False
        for line in md_lines:
            current_empty = len(line.strip()) == 0
            if not (prev_empty and current_empty):
                cleaned_lines.append(line)
            prev_empty = current_empty

        # 合并所有行
        full_text = "\n".join(cleaned_lines)

        return full_text

    def _get_elements(self) -> List:
        CRYPTO_EVAL_INFORMATION = self.CRYPTO_EVAL_INFORMATION
        DOC_MD_CONTENT = self.DOC_MD_CONTENT
        DOC_MD_TITLE = self.DOC_MD_TITLE

        # 中文到英文的映射
        title_key_map = {
            "信息种类及关键数据类型": "KEY",
            "身份鉴别": "AUTH",
            "访问控制信息完整性保护": "ACCESS_CONTROL",
            "重要数据传输机密性和完整性": "KEY_DATA_TRANSPORT_CONFIDENTIALITY_AND_INTEGRITY",
            "重要数据存储机密性": "KEY_DATA_STORE_CONFIDENTIALITY",
            "重要数据存储完整性": "KEY_DATA_STORE_INTEGRITY",
            "重要信息资源安全标记完整性": "KEY_DATA_SAFE_MARK_INTEGRITY",
            "不可否认性": "NON-REPUDIATION"
        }

        # 使用自定义配置或默认配置
        doc_key_extraction = DOC_KEY_EXTRACTION
        if self.custom_doc_key_extraction:
            # 将中文key转换为英文key
            converted_config = {}
            for chinese_key, value in self.custom_doc_key_extraction.items():
                english_key = title_key_map.get(chinese_key)
                if english_key:
                    converted_config[english_key] = value
            doc_key_extraction = converted_config

        def doc2text(filepath):
            from docx.table import _Cell, Table
            from docx.oxml.table import CT_Tbl
            from docx.oxml.text.paragraph import CT_P
            from docx.text.paragraph import Paragraph
            from docx import Document, ImagePart
            from PIL import Image
            from io import BytesIO
            import numpy as np
            from rapidocr_onnxruntime import RapidOCR
            ocr = RapidOCR()
            doc = Document(filepath)
            resp = ""

            def iter_block_items(parent):
                from docx.document import Document
                if isinstance(parent, Document):
                    parent_elm = parent.element.body
                elif isinstance(parent, _Cell):
                    parent_elm = parent._tc
                else:
                    raise ValueError("RapidOCRDocLoader parse fail")

                for child in parent_elm.iterchildren():
                    if isinstance(child, CT_P):
                        yield Paragraph(child, parent)
                    elif isinstance(child, CT_Tbl):
                        yield Table(child, parent)

            b_unit = tqdm.tqdm(total=len(doc.paragraphs) + len(doc.tables),
                               desc="RapidOCRDocLoader block index: 0")
            current_numbers = [0, 0, 0, 0, 0, 0]  # 用于跟踪各级标题编号（假设最多3级）
            CATCH_FLAG = False
            CATCH_KEY = ""
            for i, block in enumerate(iter_block_items(doc)):
                b_unit.set_description(
                    "RapidOCRDocLoader  block index: {}".format(i))
                b_unit.refresh()
                if isinstance(block, Paragraph):
                    if block.style.name.startswith("Heading") and block.text.strip() != "":
                        level = int(block.style.name.split()[-1])  # 提取标题级别（如 "Heading 1" → 1）
                        current_numbers[level - 1] += 1
                        current_numbers[level:] = [0] * (len(current_numbers) - level)
                        heading_number = ".".join(str(num) for num in current_numbers[:level] if num > 0)
                        keys = [k for k, v in doc_key_extraction.items() if v == heading_number]
                        title_text = f"{'#' * level} {heading_number} {block.text.strip()}"
                        DOC_MD_CONTENT.append(title_text)
                        DOC_MD_TITLE.append(title_text)
                        if keys:
                            CATCH_FLAG = True
                            CATCH_KEY = keys[0]
                            CRYPTO_EVAL_INFORMATION[CATCH_KEY] += title_text
                        else:
                            CATCH_FLAG = False
                        resp += f"{heading_number}{block.text.strip()}\n"
                    else:
                        paragraph_text = block.text.strip() + "\n"
                        DOC_MD_CONTENT.append(paragraph_text)
                        if CATCH_FLAG:
                            CRYPTO_EVAL_INFORMATION[CATCH_KEY] += paragraph_text
                        resp += paragraph_text
                elif isinstance(block, Table):
                    table_content = []
                    for row in block.rows:
                        row_content = []
                        for cell in row.cells:
                            # 处理单元格内的换行
                            cell_paragraphs = [p.text.strip() for p in cell.paragraphs if p.text.strip()]
                            # 使用Markdown换行符替换单元格内的换行
                            cell_text = "<br>".join(cell_paragraphs)
                            resp += cell_text + "\n"
                            row_content.append(cell_text)
                        table_content.append(row_content)
                    table_text = self.convert_table_to_md(table=table_content) + "\n"
                    DOC_MD_CONTENT.append(table_text)
                    if CATCH_FLAG:
                        CRYPTO_EVAL_INFORMATION[CATCH_KEY] += table_text
                b_unit.update(1)
            return resp

        text = doc2text(self.file_path)
        from unstructured.partition.text import partition_text
        return partition_text(text=text, **self.unstructured_kwargs)


